import baostock as bs
import pandas as pd
import db
from datetime import datetime, timedelta
import concurrent.futures
import sys
import time as time


dateFormat = "%Y-%m-%d"
cursor_stock = ""
user_id = ""

def add_days(date, days):
    now = datetime.strptime(date, dateFormat) 
    delta = timedelta(days=days)
    return datetime.strftime(now + delta, dateFormat)

def load_cursor():
    cursor_sql = "select * from `cursor` where table_name='bao_stock'"
    with db.UsingMysql() as um:
        um.cursor.execute(cursor_sql)
        cursor_rs = um.cursor.fetchall()
    return cursor_rs

def format_stock_code(stock_code):
    first = stock_code[0]
    code = stock_code
    if first == '6':
        code = "sh." + stock_code
    else:
        code = "sz." + stock_code
    return code;

def load_sharding(sharding, stock_code):
    now_str = datetime.strftime(datetime.now(), dateFormat)
    yesterday_str = add_days(now_str, -1)
    sql = "select s.* from sharding s where 1=1 "
    if sharding != None:
        sql += " and s.sharding_table=" + sharding
    if stock_code != None:
        sql += " and s.stock_code='" + stock_code + "'"
    sql += " order by stock_code asc;"
    with db.UsingMysql() as um:
        um.cursor.execute(sql);
        rs = um.cursor.fetchall()
        return rs

def load_date(sharding):
    sharding_str = "min5data" + get_sharding_str(sharding)
    sql = "select max(`date`) as `date`,stockID from " + sharding_str + " group by stockID;"
    with db.UsingMysql() as um:
        um.cursor.execute(sql)
        rs = um.cursor.fetchall()
        dic = {}
        for data in rs:
            dic[data["stockID"]] = data["date"]
        return dic

def build_sharding(cursor_rs, sharding_rs):
    dic = {}
    for data in cursor_rs:
        code = data["code"]
        dic[code] = 1
    clear_data = []
    for data in sharding_rs:
        code = data["stock_code"]
        if dic.get(code) == None:
            data["stock_code"] = code
            clear_data.append(data)
        else:
            print(code, " handled")
    return clear_data

def load_by_code(stock_code, index, start_date, end_date, update_cursor):
    load_by_day(stock_code, index, start_date, end_date, update_cursor)

def get_sharding_str(sharding):
    if(int(sharding) < 10):
        sharding = '00'+sharding
    else:
        sharding = '0'+sharding
    return sharding;


def load_by_day(stockCode, index, start, end, need_update_cursor):
    #print(stockCode+","+index)
    start_date_str = "1997-07-01"
    rs = None
    index = get_sharding_str(index)
    if start != None:
        start_date_str = start
    # else:
    #     with db.UsingMysql() as um:
    #         sql = "select max(`date`) as `date` from min5data" + index + " where stockID=%s"
    #         params = (stockCode)
    #         um.cursor.execute(sql, params)
    #         date_rs = um.cursor.fetchone()
    #         if date_rs["date"] != None:
    #             start_date = date_rs["date"]
    #             delta = timedelta(days=1)
    #             start_date = datetime.strptime(start_date, dateFormat) - delta
    #             start_date_str = datetime.strftime(start_date, dateFormat)
    # try:
    rs = query(stockCode,start_date_str, end)
    
    # except TimeoutError:
    #     load_by_day(stockCode, index)
    list = []
    i = 0
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
        date = start_date_str
        while True:
            if rs.next() == False and rs.error_code == "0":
                break
            if rs.next() == True:
                list.append({"data":rs.get_row_data(), "index": index})
                if (i + 1) % 5000 == 0:
                    begin_save(list, executor, None)
                    date = list[len(list)-1]["data"][0]
                    list = []
                i=i+1
            else:
                print(rs.error_msg)
                if rs.error_code == '10001001':
                    bs.login()
                time.sleep(1)
                rs = query(stockCode,date, end)
                
        if len(list) > 0:
            begin_save(list, executor, None)
            date = list[len(list)-1]["data"][0]
    if need_update_cursor:
        update_cursor(stockCode, date)

def query(stock_code, start_date, end_date):
    freStr = '5'
    aStr = '2'
    rs = None
    if end_date == None:
        rs = bs.query_history_k_data(format_stock_code(stock_code), "date,time,code,open,high,low,close,volume,amount",
        start_date=start_date,frequency=freStr,adjustflag=aStr)
    else:
        rs = bs.query_history_k_data(format_stock_code(stock_code), "date,time,code,open,high,low,close,volume,amount",
        start_date=start_date,frequency=freStr,adjustflag=aStr,end_date=end_date)
    return rs

def update_cursor(stock_code, date):
    select_sql = "select * from `cursor` where table_name=%s and code=%s"
    insert_sql = "insert into `cursor` (table_name, code, `value`) values(%s,%s,%s);"
    update_sql = "update `cursor` set `value`=%s where table_name=%s and code=%s"
    select_params = ("bao_stock", stock_code)
    insert_params = ('bao_stock', stock_code, date)
    update_params = (date, "bao_stock", stock_code)
    with db.UsingMysql() as um:
        um.cursor.execute(select_sql, select_params)
        data = um.cursor.fetchone()
        if data != None:
            um.cursor.execute(update_sql, update_params)
        else:
            um.cursor.execute(insert_sql, insert_params)

def begin_save(list, executor, call_back):
    return executor.submit(save_list, list)
    # save_list(list)

def save_list(list):
    stock_code = ""
    with db.UsingMysql() as um:
        check_date = ""
        need_check = True
        start_time = time.time()
        for data in list:
            if check_date != data["data"][0]:
                need_check = True
            need_check = save(data, um, need_check)
            check_date = data["data"][0]
            stock_code = data["data"][2]
        end_time = time.time()
        print('处理code:' + stock_code + '数据量：' + str(len(list)) + '  总共的时间为:', round(end_time - start_time, 3),'secs')

def save(rs, um, need_check):
    data = rs["data"]
    index = rs["index"]
    has_repeated = False
    select_params = data[1][0:12]+data[2][3:9]
    # if need_check:
    select = "select 1 from min5data" + index + " where minbasicdataID=%s"
    um.cursor.execute(select, select_params);
    selectData = um.cursor.fetchone();
    has_repeated = selectData != None
    if has_repeated:
        return True
        
    sql = "insert into min5data" + index + " (`minbasicdataID`,`date`,`datetime`,stockID,openPrice,highPrice,lowPrice,closePrice,volume,amount) "
    sql += " values('" + select_params + "','" + data[0] + "',"+ data[1] + ",'" + data[2][3:9] + "'"
    sql += ","+ data[3] + "," + data[4] + "," + data[5] + "," + data[6] + "," + data[7] + "," + data[8] + ");"
    try:
        um.cursor.execute(sql)
        um.commit()
    except Exception:
        print("exception on execute sql:" + sql)
    #     return False;
    print(data)
    return True

def load():
    cursor_rs = load_cursor()
    sharding_rs = load_sharding(None, None)
    clear_data = build_sharding(cursor_rs, sharding_rs)
    if len(clear_data) > 0:
        bs.login()
        for data in clear_data:
            load_by_code(data)   
        bs.logout()

def main():
    load()

if __name__ == '__main__':
    main()

    
